#!/bin/bash

# 错义突变微调训练启动脚本
# 基于LucaOne嵌入进行微调

# 数据集配置
DATASET_NAME="MutationPPI"
DATASET_TYPE="protein"
TASK_TYPE="binary_class"
TASK_LEVEL_TYPE="seq_level"
MODEL_TYPE="mutation_finetuning"

# 数据路径
TRAIN_DATA="../dataset/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/train/mutation_data.csv"
DEV_DATA="../dataset/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/dev/mutation_data.csv"
TEST_DATA="../dataset/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/test/mutation_data.csv"

# LucaOne嵌入路径
LUCAONE_EMBEDDINGS_DIR="../embeddings/$DATASET_NAME/lucaone/36000000"

# 输出路径
OUTPUT_DIR="../models/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/$MODEL_TYPE/$(date +%Y%m%d_%H%M%S)"

# 训练参数
BATCH_SIZE=16
LEARNING_RATE=1e-4
NUM_EPOCHS=50
MAX_LENGTH=2048
DEVICE="cuda"

# 创建输出目录
mkdir -p $OUTPUT_DIR

# 记录训练参数
cat > $OUTPUT_DIR/training_config.txt << EOF
训练配置:
- 数据集: $DATASET_NAME
- 任务类型: $TASK_TYPE
- 模型类型: $MODEL_TYPE
- 批次大小: $BATCH_SIZE
- 学习率: $LEARNING_RATE
- 训练轮数: $NUM_EPOCHS
- 最大序列长度: $MAX_LENGTH
- 设备: $DEVICE
- 输出目录: $OUTPUT_DIR
EOF

echo "开始训练错义突变微调模型..."
echo "训练数据: $TRAIN_DATA"
echo "验证数据: $DEV_DATA"
echo "测试数据: $TEST_DATA"
echo "LucaOne嵌入: $LUCAONE_EMBEDDINGS_DIR"
echo "输出目录: $OUTPUT_DIR"
echo ""

# 检查数据文件是否存在
if [ ! -f "$TRAIN_DATA" ]; then
    echo "错误: 训练数据文件不存在: $TRAIN_DATA"
    exit 1
fi

if [ ! -f "$DEV_DATA" ]; then
    echo "错误: 验证数据文件不存在: $DEV_DATA"
    exit 1
fi

if [ ! -f "$TEST_DATA" ]; then
    echo "错误: 测试数据文件不存在: $TEST_DATA"
    exit 1
fi

if [ ! -d "$LUCAONE_EMBEDDINGS_DIR" ]; then
    echo "错误: LucaOne嵌入目录不存在: $LUCAONE_EMBEDDINGS_DIR"
    exit 1
fi

# 运行训练
python train_practical_mutation.py \
    --train_data $TRAIN_DATA \
    --dev_data $DEV_DATA \
    --test_data $TEST_DATA \
    --lucaone_embeddings_dir $LUCAONE_EMBEDDINGS_DIR \
    --output_dir $OUTPUT_DIR \
    --batch_size $BATCH_SIZE \
    --learning_rate $LEARNING_RATE \
    --num_epochs $NUM_EPOCHS \
    --max_length $MAX_LENGTH \
    --device $DEVICE

# 检查训练是否成功
if [ $? -eq 0 ]; then
    echo ""
    echo "训练完成！"
    echo "模型保存在: $OUTPUT_DIR"
    echo "最佳模型: $OUTPUT_DIR/best_model.pth"
    echo "训练日志: $OUTPUT_DIR/training_config.txt"
else
    echo ""
    echo "训练失败！"
    exit 1
fi
